package com.adityayadav.mini.server;

import java.net.InetSocketAddress;
import java.net.Socket;
import java.nio.ByteBuffer;
import java.nio.channels.SelectionKey;
import java.nio.channels.Selector;
import java.nio.channels.ServerSocketChannel;
import java.nio.channels.SocketChannel;
import java.nio.channels.spi.SelectorProvider;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.LinkedBlockingQueue;

import com.adityayadav.parallelism.core.ThreadPool;

public class MiniWebSocketServer {

	
	protected ServerSocketChannel serverChannel;
	protected String host;
	protected int port;
	protected Selector selector;
	protected ThreadPool threadPool;
	protected WebSocketWorkerTask messageProcessingTask;
	protected BlockingQueue<ChangeRequest> changeRequests = new LinkedBlockingQueue<ChangeRequest>();
	protected Map pendingData = new HashMap();
	protected HashMap<SocketChannel,WebSocketWorkerTask> workers = new HashMap<SocketChannel,WebSocketWorkerTask>();

	public MiniWebSocketServer(String host, int port, ThreadPool threadPool,
			WebSocketWorkerTask messageProcessingTask) throws Exception {
		
		this.host = host;
		this.port = port;
		this.threadPool = threadPool;
		this.messageProcessingTask = messageProcessingTask;
		this.selector = initSelector();
		Thread selectorThread = new Thread() {
			public void run() {
				while (true) {
					try {
						// Process any pending changes
						while (MiniWebSocketServer.this.changeRequests.peek() != null) {
							ChangeRequest change = (ChangeRequest) MiniWebSocketServer.this.changeRequests.poll();

							switch (change.type) {
							case ChangeRequest.CHANGEOPS:
								SelectionKey key = change.socket
										.keyFor(MiniWebSocketServer.this.selector);
								key.interestOps(change.ops);
							}
						}
						

						// Wait for an event one of the registered channels
						MiniWebSocketServer.this.selector.select();

						// Iterate over the set of keys for which events are available
						Iterator selectedKeys = MiniWebSocketServer.this.selector
								.selectedKeys().iterator();
						while (selectedKeys.hasNext()) {
							SelectionKey key = (SelectionKey) selectedKeys
									.next();
							selectedKeys.remove();

							if (!key.isValid()) {
								continue;
							}

							// Check what event is available and deal with it
							if (key.isAcceptable()) {
								MiniWebSocketServer.this.accept(key);
							} else if (key.isReadable()) {
								
								WebSocketWorkerTask task = workers.get((SocketChannel) key.channel());
								if (task == null) {
									task = (WebSocketWorkerTask) MiniWebSocketServer.this.messageProcessingTask.clone();
									task.setSelectionKey(key);
									task.setServer(MiniWebSocketServer.this);
									workers.put((SocketChannel) key.channel(), task);
								}
								
								MiniWebSocketServer.this.threadPool.addTask(task);
							} else if (key.isWritable()) {
								MiniWebSocketServer.this.write(key);
							}
						}
					} catch (Exception e) {
						e.printStackTrace();
					}
				}
			}
		};
		selectorThread.start();
	}

	private Selector initSelector() throws Exception {
		// Create a new selector
		Selector socketSelector = SelectorProvider.provider().openSelector();

		// Create a new non-blocking server socket channel
		this.serverChannel = ServerSocketChannel.open();
		serverChannel.configureBlocking(false);

		// Bind the server socket to the specified address and port
		InetSocketAddress isa = new InetSocketAddress(this.host, this.port);
		serverChannel.socket().bind(isa);

		// Register the server socket channel, indicating an interest in 
		// accepting new connections
		serverChannel.register(socketSelector, SelectionKey.OP_ACCEPT);
		return socketSelector;
	}

	private void accept(SelectionKey key) throws Exception {
		// For an accept to be pending the channel must be a server socket channel.
		ServerSocketChannel serverSocketChannel = (ServerSocketChannel) key
				.channel();

		// Accept the connection and make it non-blocking
		SocketChannel socketChannel = serverSocketChannel.accept();
		Socket socket = socketChannel.socket();
		socketChannel.configureBlocking(false);

		// Register the new SocketChannel with our Selector, indicating
		// we'd like to be notified when there's data waiting to be read
		socketChannel.register(this.selector, SelectionKey.OP_READ);
	}

	public void send(SocketChannel socket, byte[] data) {
			// Indicate we want the interest ops set changed
			this.changeRequests.add(new ChangeRequest(socket,
					ChangeRequest.CHANGEOPS, SelectionKey.OP_WRITE));

			// And queue the data we want written
			synchronized (this.pendingData) {
				List queue = (List) this.pendingData.get(socket);
				if (queue == null) {
					queue = new ArrayList();
					this.pendingData.put(socket, queue);
				}
				queue.add(ByteBuffer.wrap(data));
			}
		// Finally, wake up our selecting thread so it can make the required changes
		this.selector.wakeup();
	}

	private void write(SelectionKey key) throws Exception {
		SocketChannel socketChannel = (SocketChannel) key.channel();

		synchronized (this.pendingData) {
			List queue = (List) this.pendingData.get(socketChannel);

			// Write until there's not more data ...
			while (!queue.isEmpty()) {
				ByteBuffer buf = (ByteBuffer) queue.get(0);
				try {
					socketChannel.write(buf);
				} catch (Exception e){
					closeSocketChannel(socketChannel);
				}
				if (buf.remaining() > 0) {
					// ... or the socket's buffer fills up
					break;
				}
				queue.remove(0);
			}

			if (queue.isEmpty()) {
				// We wrote away all data, so we're no longer interested
				// in writing on this socket. Switch back to waiting for
				// data.
				key.interestOps(SelectionKey.OP_READ);
			}
		}
	}

	public void closeSocketChannel(SocketChannel socketChannel) {
		removeWorker(socketChannel);
		pendingData.remove(socketChannel);
		try {
			socketChannel.close();
		} catch (Exception g){}
	}

	public void removeWorker(SocketChannel channel){
		workers.remove(channel);
	}
	class ChangeRequest {
		public static final int REGISTER = 1;
		public static final int CHANGEOPS = 2;

		public SocketChannel socket;
		public int type;
		public int ops;

		public ChangeRequest(SocketChannel socket, int type, int ops) {
			this.socket = socket;
			this.type = type;
			this.ops = ops;
		}
	}
}
